# Generated with SMOP  0.41
from smop import *
import numpy as np
from calNa import calNa
from GetImgCapacity import GetImgCapacity


# BBE.m
def BBE(im:np.ndarray,s:int):
    M,N=im.shape
    r1=np.mod(M,s)
    c1=np.mod(N,s)
    embeddImg=im
    na=calNa(s)
    p=max(np.ceil(np.log2(na)),1.0)
    capacity, BCapacity, bbBits, bIndex=GetImgCapacity(im[0:-(1+r1),0:-(1+c1)],s)
    if capacity>0:
        mess=np.concatenate(bbBits,np.random.randint(2,capacity,1))
        index=1
        for i in range(1,M-r1+1,s):
            for j in range(1,N-c1+1,s):
                bij=im[i-1:i+s-2,j-1:j+s-2]
                if bIndex[index-1]==0:
                    bij[0,0]=0
                    bij[0,1]=0
                else:
                    cb=BCapacity[i-1,j-1].cb
                    qi=BCapacity[i-1,j-1].qi
                    n0=BCapacity[i-1,j-1].n0
                    n1=BCapacity[i-1,j-1].n1
                    m=BCapacity[i-1,j-1].m
                    pos=BCapacity[i-1,j-1].pos

                    bij=bij.T
                    bij=np.ravel(bij)
                    if n0==0:
                        bij[0,1]=1
                        bij[2:-1]=mess[0:cb+1]
                        mess[0:cb+1]=np.zeros([0,cb+1])
                    elif n1==0:
                        bij[0]=1
                        bij[1]=0
                        bij[2:-1]=mess[0:cb+1]
                        mess[0:cb+1]=np.zeros([0,cb+1])
                    elif n0<n1:
                        bij[0]=0
                        bij[1:4]=1
                        pv=CalPvalue(m,p)
                        bij[2:4+p]=pv
                        posInfo=CalPosInfo(pos,qi)
                        bij[3+p:4+p+len(posInfo)]=posInfo
                        L=s**2-(4+p+len(posInfo)-1)
                        bij[4+p+len(posInfo)-1:-1]=mess[0:L+1]
                        mess=np.zeros([0,L+1])
                    else:
                        bij[0]=0
                        bij[1]=1
                        bij[2]=0
                        pv=CalPvalue(m,p)
                        bij[3:4+p]=pv
                        posInfo=CalPosInfo(pos,qi)
                        bij[4+p:4+len(posInfo)]=posInfo
                        L=s**2-(4+p+len(posInfo)-1)
                        bij[4+p+len(posInfo)-1:-1]=mess[0:L+1]
                        mess = np.zeros([0, L + 1])
                    bij=np.reshape(bij,(s,s))
                    bij=bij.T
                index+=1
                embeddImg[i-1:i+s,j-1:j+s]=bij
    return embeddImg,capacity

'''
def BBE(im=None, s=None, *args, **kwargs):
    M, N = im.shape
    # BBE.m:2
    r1 = np.mod(np.size(im, 1), s)
    # BBE.m:2
    c1 = np.mod(np.size(im, 2), s)
    # BBE.m:2
    embeddedImg = np.copy(im)
    # BBE.m:3
    na = calNa(s)
    # BBE.m:4
    p = max(np.ceil(np.log2(na)), 1)
    # BBE.m:4
    capacity, BCapacity, bbBits, bIndex = GetImgCapacity(im(np.arange(1, -(1 + r1)), np.arange(1, -(1 + c1)), s))
    # BBE.m:5
    if capacity > 0:
        mess = pd.concat([[bbBits], [np.randi(2, capacity, 1) - 1]])
        # BBE.m:7
        index = 1
        # BBE.m:8
        for i in range(1, M - r1, s).reshape(-1):
            for j in range(1, N - c1, s).reshape(-1):
                bij = im(range(i, i + s - 1), range(j, j + s - 1))
                # BBE.m:11
                if bIndex(index) == 0:
                    bij[1, 1] = 0
                    # BBE.m:13
                    bij[1, 2] = 0
                # BBE.m:13
                else:
                    cb = BCapacity(i, j).cb
                    # BBE.m:15
                    qi = BCapacity(i, j).qi
                    # BBE.m:16
                    n0 = BCapacity(i, j).n0
                    # BBE.m:17
                    n1 = BCapacity(i, j).n1
                    # BBE.m:18
                    m = BCapacity(i, j).m
                    # BBE.m:19
                    pos = BCapacity(i, j).pos
                    # BBE.m:20
                    bij = bij.T
                    # BBE.m:22
                    bij = np.ravel(bij)
                    # BBE.m:22
                    if n0 == 0:
                        bij[np.arange(1, 2)] = 1
                        # BBE.m:24
                        bij[np.arange(3, -1)] = mess(np.arange(1, cb))
                        # BBE.m:24
                        mess[np.arange(1, cb)] = []
                    # BBE.m:24
                    else:
                        if n1 == 0:
                            bij[1] = 1
                            # BBE.m:26
                            bij[2] = 0
                            # BBE.m:26
                            bij[np.arange(3, -1)] = mess(np.arange(1, cb))
                            # BBE.m:26
                            mess[np.arange(1, cb)] = []
                        # BBE.m:26
                        else:
                            if n0 < n1:
                                bij[1] = 0
                                # BBE.m:28
                                bij[np.arange(2, 3)] = 1
                                # BBE.m:28
                                pv = CalPvalue(m, p)
                                # BBE.m:29
                                bij[np.arange(4, 4 + p - 1)] = pv
                                # BBE.m:29
                                posInfo = CalPosInfo(pos, qi)
                                # BBE.m:30
                                bij[np.arange(4 + p, 4 + p + len(posInfo) - 1)] = posInfo
                                # BBE.m:30
                                L = s ** 2 - (4 + p + len(posInfo) - 1)
                                # BBE.m:31
                                bij[np.arange(4 + p + len(posInfo), -1)] = mess(np.arange(1, L))
                                # BBE.m:32
                                mess[np.arange(1, L)] = []
                            # BBE.m:32
                            else:
                                bij[1] = 0
                                # BBE.m:34
                                bij[2] = 1
                                # BBE.m:34
                                bij[3] = 0
                                # BBE.m:34
                                pv = CalPvalue(m, p)
                                # BBE.m:35
                                bij[np.arange(4, 4 + p - 1)] = pv
                                # BBE.m:35
                                posInfo = CalPosInfo(pos, qi)
                                # BBE.m:36
                                bij[np.arange(4 + p, 4 + p + len(posInfo) - 1)] = posInfo
                                # BBE.m:36
                                L = s ** 2 - (4 + p + len(posInfo) - 1)
                                # BBE.m:37
                                bij[np.arange(4 + p + len(posInfo), -1)] = mess(np.arange(1, L))
                                # BBE.m:38
                                mess[np.arange(1, L)] = []
                    # BBE.m:38
                    bij = np.reshape(bij, pd.concat([s, s]))
                    # BBE.m:40
                    bij = bij.T
                # BBE.m:40
                index = index + 1
                # BBE.m:42
                embeddedImg[np.arange(i, i + s - 1), np.arange(j, j + s - 1)] = bij
    # BBE.m:43

    return embeddedImg, capacity

'''


# pv数组
'''
def CalPvalue(p,L):
     pv=np.base_repr(p,2,L)
     pv=[int(x) for x in str(pv)]
     if len(pv)>L:
         pv=pv[1:-1]
    return pv
'''


def CalPvalue(p, L):
    pv = np.base_repr(p, 2, L)
    pv = [int(x) for x in str(pv)]
    if len(pv) > L:
        pv = pv[1:-1]
    return pv


'''
def CalPvalue(p=None, L=None, *args, **kwargs):


    #pv = str2num(dec2bin(p, L).T)# wrong
    pv=np.base_repr(p,2,L)
    pv=[int(x) for x in str(pv)]
    # BBE.m:51
    if len(pv) > L:
        pv = pv(np.arange(2, -1))
    # BBE.m:53

    return pv
def CalPosInfo(pos=None, qi=None, *args, **kwargs):


    posInfo = []
    # BBE.m:58
    for i in range(1, len(pos)).reshape(-1):
        if i == 1:
            posInfo = pd.concat([[posInfo], [CalPvalue(pos(i), qi(i))]])
        # BBE.m:61
        else:
            posInfo = pd.concat([[posInfo], [CalPvalue(pos(i) - pos(i - 1), qi(i))]])
    # BBE.m:63

    return posInfo

'''


def CalPosInfo(pos, qi):
    posInfo = np.asarray([])
    for i in range(1, len(pos) + 1):
        if i == 1:
            posInfo = np.concatenate(posInfo, np.asarray(CalPvalue(pos(i - 1), qi(i))))
        else:
            posInfo = np.concatenate(posInfo, np.asarray(CalPvalue(pos(i - 1) - pos(i - 2), qi(i))))
    return posInfo
